package com.zhm.demo.nats;

import java.net.*;
import java.nio.ByteBuffer;
import java.security.SecureRandom;
import java.util.Arrays;

public class StunClient {
    private static final String STUN_SERVER = "stun.l.google.com";
    private static final int STUN_PORT = 19302;

    // STUN 消息类型（Binding Request）
    private static final byte[] BINDING_REQUEST = {
            0x00, 0x01,             // Message Type: Binding Request (0x0001)
            0x00, 0x00,             // Message Length
            0x21, 0x12, 0x04, 0x42  // Magic Cookie (固定值)
    };

    public static void main(String[] args) throws Exception {
        // 创建 UDP Socket
        DatagramSocket socket = new DatagramSocket();
        socket.setSoTimeout(5000);  // 超时 5 秒

        // 构建 STUN 请求报文（含随机 Transaction ID）
        byte[] transactionId = new byte[12];
        new SecureRandom().nextBytes(transactionId);  // 随机生成 12 字节事务ID
        ByteBuffer buffer = ByteBuffer.allocate(20);
        buffer.put(BINDING_REQUEST);
        buffer.put(transactionId);
        byte[] requestData = buffer.array();

        // 发送请求到 STUN 服务器
        InetAddress stunAddress = InetAddress.getByName(STUN_SERVER);
        DatagramPacket request = new DatagramPacket(requestData, requestData.length, stunAddress, STUN_PORT);
        socket.send(request);

        // 接收响应
        byte[] responseData = new byte[512];
        DatagramPacket response = new DatagramPacket(responseData, responseData.length);
        socket.receive(response);

        // 解析响应中的 XOR-MAPPED-ADDRESS
        parseStunResponse(response.getData(), response.getLength());
        socket.close();
    }

    private static void parseStunResponse(byte[] data, int length) {
        if (length < 20) {
            System.err.println("Invalid STUN response");
            return;
        }

        // 检查 Magic Cookie
        byte[] magicCookie = Arrays.copyOfRange(data, 4, 8);
        if (!Arrays.equals(magicCookie, new byte[]{0x21, 0x12, (byte) 0xA4, 0x42})) {
            System.err.println("Invalid Magic Cookie");
            return;
        }

        // 遍历属性列表
        int offset = 20;  // 跳过头部
        while (offset < length) {
            int attrType = ((data[offset] & 0xFF) << 8) | (data[offset + 1] & 0xFF);
            int attrLength = ((data[offset + 2] & 0xFF) << 8) | (data[offset + 3] & 0xFF);
            offset += 4;

            if (attrType == 0x0020) {  // XOR-MAPPED-ADDRESS 类型
                // 解析 IP 和端口（异或解码）
                int port = ((data[offset + 2] & 0xFF) << 8) | (data[offset + 3] & 0xFF);
                port ^= 0x2112;  // XOR 解码

                int ip = ByteBuffer.wrap(data, offset + 4, 4).getInt();
                ip ^= 0x2112A442;  // XOR 解码 Magic Cookie

                String publicIp = String.format("%d.%d.%d.%d",
                        (ip >> 24) & 0xFF, (ip >> 16) & 0xFF, (ip >> 8) & 0xFF, ip & 0xFF);
                System.out.println("公网地址: " + publicIp + ":" + port);
                return;
            }
            offset += attrLength;
        }
    }
}
